/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Declarations of ray classes.

#ifndef __ray__
#define __ray__

#include "config.h"
#include "mcrx-types.h"
#include "blitz/array.h"
#include "mono_poly_abstract.h"

#ifdef HAVE_BOOST_SERIALIZATION
#include <boost/serialization/serialization.hpp>
#include <boost/serialization/base_object.hpp>
#include <boost/serialization/split_member.hpp>
#endif

namespace mcrx {
  class ray_base;
  template <typename> class ray;
  template <typename T_numtype>
  class ray<blitz::Array<T_numtype, 1> >;

  template<typename T> 
  ray<T> independent_copy(const ray<T>& rhs);
}

/** A ray_base is a base class for rays, containing basic
    functionality that has to do with positions, directions, etc, that
    does not depend on what the ray carries. */
class mcrx::ray_base {
private:
  vec3d dir_, idir_; ///< Direction inverse, precomputed to save time.
  vec3d pos_;
  T_float l_; ///< Path length traversed.
  T_float d_; ///< Doppler shift factor (used when doing kinematic calculations).
  int nscat_; ///< Number of scatterings.
  T_densities n_; ///< Array of column densities of species traversed by ray.

  void assert_stuff() {
    assert (dot (dir_, dir_) > .9999); assert (dot (dir_, dir_)< 1.0001);
  };

#ifdef HAVE_BOOST_SERIALIZATION
  friend class boost::serialization::access;
  /** For efficiency, we only serialize the length and data of the
      column density member instead of the full array object. */
  template<class T_arch>
  void save(T_arch& ar, const unsigned int version) const {
    ar << dir_ << idir_ << pos_ << l_ << d_ << nscat_;
    const boost::serialization::collection_size_type nsize(n_.size());
    ar << BOOST_SERIALIZATION_NVP(nsize);
    assert(n_.stride(blitz::firstDim)==1 || n_.size()==0);
    if(n_.size()>0)
      ar << boost::serialization::make_array(n_.dataFirst(), nsize);

  };
  template<class T_arch>
  void load(T_arch& ar, const unsigned int version) {
    ar >> dir_ >> idir_ >> pos_ >> l_ >> d_ >> nscat_;
    boost::serialization::collection_size_type nsize;
    ar >> BOOST_SERIALIZATION_NVP(nsize);
    n_.resize(nsize);
    if(nsize>0)
      ar >> boost::serialization::make_array(n_.dataFirst(), nsize);
  };
  BOOST_SERIALIZATION_SPLIT_MEMBER()
#endif

public:
  ray_base() {};
  ray_base (const vec3d& pos, const vec3d& dir):
    dir_ (dir), idir_ (1./dir), pos_ (pos), l_ (0), d_(1.0), nscat_ (0) {
    assert_stuff();};

  /** Constructor makes a reference copy of the column densities. */
  ray_base (const vec3d& pos, const vec3d& dir, const T_densities& nn):
    dir_ (dir), idir_ (1./dir), pos_ (pos), 
    n_(nn), l_ (0), d_(1.0), nscat_ (0) { assert_stuff();};
  /** Copy constructor makes a reference copy of the column densities. */
  ray_base(const ray_base& rhs) : 
    dir_(rhs.dir_), idir_(rhs.idir_), pos_(rhs.pos_), 
    n_(rhs.n_), l_(rhs.l_), d_(rhs.d_), nscat_(rhs.nscat_) {};

  /** Assignment operator does NOT make sure the column density array
      is the right size. If rhs has an uninitialized column density,
      it just ignores it. */
  ray_base& operator= (const ray_base& rhs) {
    if (this == &rhs) return *this;
    dir_ = rhs.dir_; idir_ = rhs.idir_; pos_ = rhs.pos_; l_ = rhs.l_; 
    d_ = rhs.d_; nscat_ = rhs.nscat_; 
    if(rhs.n_.size()>0) n_ = rhs.n_; else n_=0;
    return *this;
  };

  /** A "deep assignment" operator that makes this an independent copy
      of rhs. */
  ray_base& deep_assign(const ray_base& rhs) {
    if (this == &rhs) return *this;
    dir_ = rhs.dir_; idir_ = rhs.idir_; pos_ = rhs.pos_; l_ = rhs.l_; 
    d_ = rhs.d_; nscat_ = rhs.nscat_; 
    n_.reference(rhs.n_.copy());
    return *this;
  };

  /** A "shallow assignment" operator that makes this a copy of rhs
      that both reference the same column density vector. */
  ray_base& shallow_assign(const ray_base& rhs) {
    if (this == &rhs) return *this;
    dir_ = rhs.dir_; idir_ = rhs.idir_; pos_ = rhs.pos_; l_ = rhs.l_; 
    d_ = rhs.d_; nscat_ = rhs.nscat_; 
    n_.reference(rhs.n_);
    return *this;
  };

  const vec3d& position () const {return pos_;};
  const vec3d& direction () const {return dir_;};
  const vec3d& inverse_direction () const {return idir_;};
  const T_densities& traversed_column() const {return n_;};
  T_float length () const {return l_;};
  T_float doppler () const {return d_;};
  /// Set the ray doppler factor to d.
  void set_doppler (T_float d) {assert((d>0.1)&&(d<2)); d_=d;};
  /// Add a doppler shift of d to the ray doppler factor.
  void add_doppler (T_float d) {assert((d>0.1)&&(d<2)); d_*=d; 
    assert((d_>0.1)&&(d_<2));};

  int scatterings () const {return nscat_;};
  int& scatterings () {return nscat_;};

  // Note that setting the ray position resets traversed length and column.
  void set_position (const vec3d& pos) {pos_ = pos; l_ = 0; n_ = 0;};
  void set_direction (const vec3d& dir) {
    dir_ = dir; idir_ = 1./dir; assert_stuff(); };

  void set_traversed_column (const T_densities& n) {n_ = n; };
  void reset_traversed_column () {n_ = 0; };
  void reset_length () {l_ = 0; };
  /** Sets the column density to reference another array. */
  void set_column_ref(T_densities& n) { n_.reference(n); };

  /** Scatters the ray in the direction defined by cos_theta and phi
      (in the ray frame). */
  void scatter_in_direction (T_float cos_theta, T_float phi);

  /// Propagates the ray a specified length, ignoring column density
  void propagate (T_float dl) {
    pos_ += dl*direction (); l_ += dl;};
  /// Propagates the ray a specified length and optical depth.
  void propagate (T_float dl, const T_densities& dn) {
    pos_ += dl*direction (); l_ += dl; 
    if(n_.size()==1) n_(0) += dn(0); else n_ += dn;};
  template<typename T>
  void propagate (T_float dl, blitz::_bz_ArrayExpr<T> expr) {
    pos_ += dl*direction (); l_ += dl; n_ += expr;};

  /** Checks if the ray intersects an axis-aligned cube between
      coordinates min and max and returns the length to
      intersection. NaN is returned if the ray misses the cube. If the
      ray is on the entry face and allow_zero is returned, 0 will be
      returned, otherwise the exit face distance will be returned. */
  T_float intersect_cube(const vec3d& min, const vec3d& max, bool assume_inside) const;
};


/** The basic ray template where the intensity is a scalar. */
template <typename lambda_type>
class mcrx::ray: public ray_base {
public:
  typedef lambda_type T_lambda;
private: 
  T_lambda i_; ///< The intensity of the ray

#ifdef HAVE_BOOST_SERIALIZATION
  /// Serialization support.
  friend class boost::serialization::access;
  template<class T_arch>
  void serialize(T_arch& ar, const unsigned int version) {
    ar & boost::serialization::base_object<ray_base>(*this);
    ar & i_;
  };
#endif

public:
  ray() {};
  ray (const vec3d& p, const vec3d& d, const T_densities& nn, 
       const T_lambda& ii):
    ray_base (p, d, nn), i_ (ii) {};
  ray (const vec3d& p, const vec3d& d, const T_lambda& ii):
    ray_base (p, d), i_ (ii) {};

  /** A "shallow assignment" operator that makes this a copy of rhs
      that both reference the same column density vector. */
  ray& shallow_assign(const ray& rhs) {
    if (this == &rhs) return *this;
    ray_base::shallow_assign(rhs);
    i_ = rhs.i_;
    return *this;
  };

  const T_lambda& intensity () const {return i_;};
  void set_intensity (T_lambda i) {i_ = i; };
  /** For these rays, set_intensity_ref just copies it. */
  void set_intensity_ref(const T_lambda& i) { i_ = i; };
  void bias(const T_lambda& factor) {
    assert(condition_all(factor>=0));
    assert(condition_all(factor<blitz::huge(T_float())));
    ;i_ *= factor;};
};


/** Specialization of the ray template for the polychromatic case when
    the intensity is a blitz array. */
template <typename T_numtype>
class mcrx::ray<blitz::Array<T_numtype, 1> >: public mcrx::ray_base {
public:
  typedef typename blitz::Array<T_numtype, 1> T_lambda;
private: 
  T_lambda i_; ///< The intensity of the ray

#ifdef HAVE_BOOST_SERIALIZATION
  /// Serialization support.
  friend class boost::serialization::access;
  template<class T_arch>
  void save(T_arch& ar, const unsigned int version) const {
    ar << boost::serialization::base_object<ray_base>(*this);
    const boost::serialization::collection_size_type isize(i_.size());
    ar << BOOST_SERIALIZATION_NVP(isize);
    assert(i_.stride(blitz::firstDim)==1 || isize==0);
    if(isize>0)
     ar << boost::serialization::make_array(i_.dataFirst(), isize);
  };
  template<class T_arch>
  void load(T_arch& ar, const unsigned int version) {
    ar >> boost::serialization::base_object<ray_base>(*this);
    boost::serialization::collection_size_type isize;
    ar >> BOOST_SERIALIZATION_NVP(isize);
    i_.resize(isize);
    assert(i_.stride(blitz::firstDim)==1 || isize==0);
    if(isize>0)
      ar >> boost::serialization::make_array(i_.dataFirst(), isize);
  };
  BOOST_SERIALIZATION_SPLIT_MEMBER()
#endif

public:
  ray () {};
  /** Constructor makes a reference copy of the intensity array. */
  ray (const vec3d& p, const vec3d& d, const T_densities& nn, 
       const T_lambda& ii):
    ray_base (p, d, nn), i_ (ii) {};
  /** Constructor makes a reference copy of the intensity array. */
  ray (const vec3d& p, const vec3d& d, const T_lambda& ii):
    ray_base (p, d), i_(ii) {};
  /** Copy constructor makes a reference copy of the intensity array. */
  ray (const ray& rhs) : ray_base(rhs), i_(rhs.i_) {};

  /** Assignment operator. Does NOT make sure the intensity array is
      the right size. */
  ray& operator= (const ray& rhs) {
    if (this == &rhs) return *this;
    ray_base::operator= (rhs);
    i_ = rhs.i_;
    return *this;
  };

  /** A "deep assignment" operator that makes this an independent copy
      of rhs. */
  ray& deep_assign(const ray& rhs) {
    if (this == &rhs) return *this;
    ray_base::deep_assign(rhs);
    i_.reference(rhs.i_.copy());
    return *this;
  };

  /** A "shallow assignment" operator that makes this a copy of rhs
      that both reference the same column density vector. */
  ray& shallow_assign(const ray& rhs) {
    if (this == &rhs) return *this;
    ray_base::shallow_assign(rhs);
    i_.reference(rhs.i_);
    return *this;
  };

  const T_lambda& intensity () const {return i_;};
  void set_intensity (T_float i) {i_ = i;};
  /// Sets the intensity, which must be appropriately sized first.
  void set_intensity (const T_lambda& i) {i_ = i; };
  template<typename T>
  void set_intensity (blitz::_bz_ArrayExpr<T> i) {i_= i;};

  /** Sets the intensity to refer to another array. */
  void set_intensity_ref(const T_lambda& i) { i_.reference(i); };

  void bias(T_float b) {
    assert(b>=0);
    assert(b<blitz::huge(T_float()));
    i_ *= b;};
  void bias(const T_lambda& b) {
    assert(condition_all(b >= 0));
    assert(condition_all(b < blitz::huge(T_float())));
    i_ *= b;};
  template<typename T>
  void bias(blitz::_bz_ArrayExpr<T> b) {
    assert(condition_all(b >= 0));
    assert(condition_all(b < blitz::huge(T_float())));
    i_ *= b;};

  template <typename T_expr>
  void effective_scatter(T_lambda& norm, const T_lambda& lambda, 
			 const blitz::ETBase<T_expr>& em,
			 T_float cos_theta, T_float phi);
			 
};

/** Makes a deep copy of a ray. This template should work for all ray
    types. */
template<typename T>
mcrx::ray<T>
mcrx::independent_copy(const ray<T>& rhs)
{
  ray<T> r(rhs);
  
  T_densities ncopy(independent_copy(rhs.traversed_column()));
  r.set_column_ref(ncopy);
  r.set_intensity_ref(independent_copy(rhs.intensity()));
  return r;
}

/** Performs an "effective scattering", changing the spectrum of the
    ray while conserving energy. Emission is the outgoing spectrum of
    the ray, normalized to unit luminosity. The direction to scatter
    the ray in is specified by cos_theta and phi, which is given in
    the LAB frame. This is called by the effective_scatter method in
    the emission classes. */
template <typename T_numtype>
template <typename T_expr>
void
mcrx::ray<blitz::Array<T_numtype, 1> >::
effective_scatter(T_lambda& norm, const T_lambda& lambda, 
		  const blitz::ETBase<T_expr>& em,
		  T_float cos_theta, T_float phi)
{
  const T_expr& emission(em.unwrap());

  // we want to keep some notion of the original ray intensity, so
  // that less intense rays become more likely to be RR'd away as they
  // go on. For this reason, we set the ray intensity to the max
  // incoming intensity. If we have no emission yet, we just retain
  // the incoming intensity.
  if(all(emission==emission)) {
    const T_float current_lum = integrate_quantity(intensity()*norm, 
						   lambda, false);
    const T_float current_norm = integrate_quantity(norm, 
						    lambda, false);
    assert(current_lum>1e-30);

    // we set the new norm to have the shape of the emission and the
    // same integral as the old norm.
    norm = current_norm*emission;

    // the (constant) intensity is then determined since we need to
    // get the correct luminosity
    set_intensity(current_lum/current_norm);

    // now we apply the same dynamic_range procedure to lower the
    // intensity in regions where the normalization is small
    const T_float sed_dynamic_range=1e-3;
    if (sed_dynamic_range>0) {
      const T_float maxnorm = max(norm);
      bias(where(norm < sed_dynamic_range*maxnorm,
		 (1./(maxnorm*sed_dynamic_range))*norm, 1.0));
      // norm==0 gives a NaN here unless we deal with it separately
      norm *= where(norm < sed_dynamic_range*maxnorm && norm>0,
		    (maxnorm*sed_dynamic_range)/norm, 1.0);
      DEBUG(1,								\
	    const T_float postlum = integrate_quantity(intensity()*norm, \
						       lambda, false);	\
	    const T_float postnorm = integrate_quantity(norm, lambda, false); \
	    assert(abs(postlum/current_lum-1)<1e-10);			\
	    assert((postnorm>0.99*current_norm) && (postnorm<2*current_norm)); \
	    ASSERT_ALL(norm < 1e60); );
    }
  }

  const T_float sin_theta = sqrt (1- cos_theta*cos_theta);
  const vec3d new_dir(sin_theta*cos (phi), sin_theta*sin (phi), 
		      cos_theta);
  set_direction(new_dir);
  reset_traversed_column();
  ++scatterings();
}


#endif
